#!/bin/sh
# - mode
# - nr_dev
# - nr_part
# - part_type
# - ns_id
# - ns_mode

. "$LKP_SRC"/lib/nvdimm.sh

: "${mode:=dax}"
: "${part_type:=-}"

keep_quiet()
{
	[ -n "$quiet_mode" ]
}

__probe_devs()
{
	case "$mode" in
		raw|fsdax)
			rdevs=$(ls /dev/pmem* | grep '/dev/pmem[0-9]\+$')
			;;
		dax)
			rdevs=$(ls /dev/dax*)
			ignore_parts=1
			;;
		*)
			echo "Invalid nvdimm mode: $mode!" >&2
			exit 1
			;;
	esac
	rnr_dev=$(echo $rdevs | wc -w)
	: "${nr_dev:=$rnr_dev}"

	if [ "$rnr_dev" -eq 0 ]; then
		keep_quiet || echo "No $mode devices found!" >&2
		return 1
	fi

	if [ "$rnr_dev" -lt "$nr_dev" ]; then
		keep_quiet || echo "No enough $mode devices found: asked $nr_dev, got $rnr_dev" >&2
		return 1
	fi

	devs=$(echo $rdevs | cut "-f1-$nr_dev" -d' ')
}

probe_devs()
{
	quiet_mode=1
	for i in $(seq 5); do
		__probe_devs && return
		sleep 5
	done

	quite_mode=
	__probe_devs
}

probe_parts()
{
	rparts=
	for dev in $devs; do
		rparts_d=$(ls "$dev"* | grep "$dev.\+")
		rparts="$rparts $rparts_d"
	done

	rnr_part=$(echo $rparts | wc -w)

	if [ -z "$nr_part" ]; then
		if [ "$rnr_part" -ne 0 ]; then
			keep_quiet || echo "Cannot delete all partitions: $rparts" >&2
			return 1
		fi
		parts=$devs
		nr_part=$nr_dev
		return 0
	fi

	if [ "$rnr_part" -ne "$nr_part" ]; then
		keep_quiet || echo "Partition number not match: asked $nr_part, got $rnr_part: $rparts" >&2
		return 1
	fi

	# normalize whitespace
	parts=$(echo $rparts)
}

create_parts()
{
	if [ -z "$nr_part" ]; then
		for dev in $devs; do
			/sbin/sfdisk --delete $dev
		done
		return
	fi

	nr_part_d=$((nr_part / nr_dev))
	for dev in $devs; do
		bdev=$(basename "$dev")
		# size in block (512)
		size_d=$(cat /sys/block/"$bdev"/size)
		# avoid to use 1M at the begin of the device
		size_p=$(((size_d - 2048) / nr_part_d / 2 / 1024))
		(
			echo "label: gpt"
			for i in $(seq $nr_part_d); do
				echo "- ${size_p}M $part_type"
			done
		) | /sbin/sfdisk $dev
	done
}

probe_create_parts()
{
	quiet_mode=1
	for i in $(seq 5); do
		probe_parts && return
		sleep 1
	done
	create_parts
	for i in $(seq 5); do
		probe_parts && return
		sleep 5
	done
	quiet_mode=
	probe_parts
}

configure_nvdimm

probe_devs

if [ -n "$ignore_parts" ]; then
	parts=$devs
	nr_part=$nr_dev
else
	probe_create_parts
fi

cat >> $TMP/env.yaml <<EOF

# setup/nvdimm
nr_partitions: $nr_part
partitions: $parts
disk_description: NVDIMM devices
EOF
